#include "BugTracker.h"

BugTracker::BugTracker()
{
	bIsInit = false;
}

BugTracker::~BugTracker()
{
	for( int k = 0; k < numcams; k++ ) {
		cvReleaseMat( &(realViews[k]) );
		cvReleaseMat( &(genViews[k]) );
	}
		
}

int ssw = -1;
void BugTracker::Init( int numCams, const valarray<float>& initT, const valarray<float>& initR, const valarray<float>& initAng )
{
	// initialize: true parameters known
	this->prevR			= initR; 
	this->prevT			= initT; 
	this->prevAng		= initAng; 
	
	this->init_predR	= initR; 
	this->init_predT	= initT;	
	this->init_predAng	= initAng;
	
	this->curR			= initR;
	this->curT			= initT;
	this->curAng		= initAng;

	this->predR			= initR; 
	this->predT			= initT; 
	this->predAng		= initAng; 

	this->realViews		= vector<CvMat*>(numCams);
	this->genViews		= vector<CvMat*>(numCams);
	this->numcams		= numCams;

	for( int i =0; i < numCams; i++ ) {
		realViews[i] = NULL;
		genViews[i]  = NULL;
	}
	steps = 0;
	this->bIsConverged	= true;
	this->bIsInit		= true;

	this->pixelIdxReal = vector< vector<int>  >(numCams);
	this->pixelIdxGen  = vector< vector<int>  >(numCams);
	this->pixelIdxGenPrev  = vector< vector<int>  >(numCams);
	this->pixelIdxGenIn  = vector< vector<int> >(numCams);
	this->pixelIdxGenOut  = vector< vector<int>  >(numCams);

	
	 distBestUnNormalized = vector<float>(numCams);
	 distCurUnNormalized = vector<float>(numCams);
	segPixCount = vector<int>(numCams);
	
	attempts = 0;
	distBest = 1e7;

}

float BugTracker::GetDistModHausdorff(int camid)
{
	float dist = this->distBestUnNormalized[camid];
	int step = 1;
	// the 'synthetic' image is binary, trivially segmentable
	// the 'real' one is (slightly?) harder to segment


	//for( int camid = 0; camid < this->numcams; camid++ ) {
		int imgH = this->genViews[camid]->height;
		int imgW = this->genViews[camid]->width;
		vector<int> inIdx = this->pixelIdxGenIn[camid];
		vector<int> outIdx = this->pixelIdxGenOut[camid];
		vector<int> realIdx = this->pixelIdxReal[camid];
		for( int k = 0; k < inIdx.size(); k+=step ) {
			if( k >= inIdx.size() )
				continue;
			float minDist = 1e7;
			int i = inIdx[k] / imgW; // y-coordinate Omega1
			int j = inIdx[k] % imgW; // x-coordinate Omega1
			for( int m = 0; m < realIdx.size(); m++ ) { // closest point in Omega2 to this k-th point in Omega1
				int ir = realIdx[m] / imgW;// y-coordinate Omega2
				int jr = realIdx[m] % imgW;// x-coordinate Omega2
				float ds = sqrt( float( (ir-i)*(ir-i) + (jr-i)*(jr-j) ) );
				if( ds < minDist ) {
					minDist = ds;
					if( ds < 1e-12 )
						break;
				}
					
			}
			dist += minDist;
		}
		for( int k = 0; k < outIdx.size(); k++ ) {
			float minDist = 1e7;
			int i = outIdx[k] / imgW; // y-coordinate Omega1
			int j = outIdx[k] % imgW; // x-coordinate Omega1
			for( int m = 0; m < realIdx.size(); m++ ) { // closest point in Omega2 to this k-th point in Omega1
				int ir = realIdx[m] / imgW;// y-coordinate Omega2
				int jr = realIdx[m] % imgW;// x-coordinate Omega2
				float ds = sqrt( float( (ir-i)*(ir-i) + (jr-i)*(jr-j) ) );
				if( ds < minDist ) {
					minDist = ds;
					if( ds < 1e-12 )
						break;
				}
			}
			dist -= minDist;
		}
		if( dist < 0 ) {
			cout<<"doh!";
			int breakhere = 1;}
//	}
	return dist;
}


void BugTracker::SetOutFile(fstream* file) {
	fileout = file;
}

void BugTracker::SaveToFile() {
	// now get est. params
	(*fileout)<<steps<<","<<curR[0]<<","<<curR[1]<<","<<curR[2];
	(*fileout)<<","<<curT[0]<<","<<curT[1]<<","<<curT[2];
	for( int i = 0; i < curAng.size(); i++ ) {
		(*fileout)<<","<<curAng[i];
	}
	(*fileout)<<"\n";
}

void BugTracker::GetBestEstParams( valarray<float>& T, valarray<float>& R, valarray<float>& Ang )
{
	T = curT;
	R = curR;
	Ang = curAng;

}

void BugTracker::LinearPredictNextParams()
{
	init_predT = curT + (curT - prevT);
	init_predR = curR + (curR - prevR); 
	init_predAng = curAng + (curAng - prevAng);

}

void BugTracker::UpdateEstParams( valarray<float>& T, valarray<float>& R, valarray<float>& Ang ) {
	// using one of possibly several techniques, update the parameters and send them back to the renderer

	// note: if the frame we just got minimizes the distance, we had the right params last time
	// so return the 'current' value

	distCur = 0.0;
	for( int k = 0; k < this->numcams; k++ ) {
		this->distCurUnNormalized[k] = this->GetDistModHausdorff(k);
		this->distCur += this->distCurUnNormalized[k] / this->segPixCount[k];
	}
	int perMode = 300;
	int totalIts = (2+6*2)*perMode;

	if(distCur < 1e-3 || totalIts == attempts) { // convergence criteria on previous value
		cout<<"satisfied convergence with error "<<distCur<<", iterations "<<attempts<<", for step "<<this->steps<<"\n";
		bIsConverged = true;
		T = curT;
		R = curR;
		Ang = curAng;

		// form the next iteration's predicted values
		this->LinearPredictNextParams();
	
		// current -> previous
		prevT = curT;
		prevR = curR;
		prevAng = curAng;
		
		SaveToFile();

		return;
	}
	 
	/* get the distance generated between the current
	 real image and the 'generated' one from sending params last time */
	

	int nextMode = attempts / perMode;
	if( 0 == nextMode )
		paramMode = 0;
	if( 1 == nextMode )
		paramMode = 2;
	if( 2 == nextMode )
		paramMode = 3;
	if( 3 == nextMode )
		paramMode = 4;
	if( 4 == nextMode )
		paramMode = 5;
	if( 5 == nextMode )
		paramMode = 6;
	if( 6 == nextMode )
		paramMode = 7;
	if( 7 == nextMode )
		paramMode = 1;
	if( 8 == nextMode )
		paramMode = 2;
	if( 9 == nextMode )
		paramMode = 3;
	if( 10 == nextMode )
		paramMode = 4;
	if( 11 == nextMode )
		paramMode = 5;
	if( 12 == nextMode )
		paramMode = 6;
	if( 13 == nextMode )
		paramMode = 7;
	cout<<"parammode = "<<paramMode<<"\n";
	
	/*int rigidits = 80;
	int perAngGroup = 40;
	if(  rigidits == attempts) 
		paramMode = 2;
	if( paramMode != 1 && attempts % perAngGroup == 0 && attempts > rigidits ) {
		paramMode++;
		if(paraMode == 8 )
			paramMode = 1;
	}*/

	if(distCur < distBest) { // if the image we just got shrunk the distance
		cout<<"improved with "<<distCur<<"\n";
		curT	= predT;
		curR	= predR;
		curAng	= predAng;
	
		distBest = distCur;
		distBestUnNormalized = distCurUnNormalized;
		for( int k = 0; k < this->numcams; k++ ) 	{
			pixelIdxGenPrev[k] = pixelIdxGen[k];
			int brteakhere = 1;
		}
	

	}
	this->StochasticDescentRigid(); // update R,T
	// create predT, predR, predAng from modifications of current values
	/*if( attempts >= 2) {
			curT = truT;
			curR = truR;
			curAng = truAng;
			predT = truT;
			predR = truR;
			predAng = truAng;
		}*/

	// a predicted value to return
	T = predT;
	R = predR;
	Ang = predAng;
	attempts++;
	srand( time(NULL) );

	
}

void BugTracker::StochasticDescentRigid()
{
	// update predictions, stochastic descent from current values
	/*predT[0]		=curT[0]	 ;
	predT[1]		=curT[1] + 2*rand()%1000 * 1e-3	-1.0 ;
	predT[2]		=curT[2] + 2*rand()%1000 * 1e-3	-1.0 ;

	predR[0]		=curR[0] + 2*rand()%1000 * 1e-3	-1.0 ;
	predR[1]		=curR[1]	 ;
	predR[2]		=curR[2]	 ;*/

	
	
	
	std::tr1::mt19937 eng;  // a core engine class ('mersenne twister')
	float sigma = 0.5;
	if( paramMode > 1 )
		sigma = 0.1;
	std::tr1::normal_distribution<float> gen1(0, sigma);
	eng.seed( time(NULL) );
	float d1 = gen1(eng);
	float d2 = gen1(eng);
	float d3 = gen1(eng); 
	predT = curT;
	predR = curR;
	if( 0 == paramMode || 1 == paramMode ) { // rigid body params
		predT[0]		=init_predT[0] ;
		predT[1]		=init_predT[1] + 0*d1 ;
		predT[2]		=init_predT[2] + d2 ;
		predR[0]		=init_predR[0] + 3*d3 ;
		predR[1]		=init_predR[1]	 ;
		predR[2]		=init_predR[2]	 ;
		predAng = curAng;
		init_predAng = curAng;
	}
	//predAng = curAng;
	if( 2 <= paramMode ) {
		init_predT = curT;
		init_predR = curR;
		init_predAng = curAng;
		//for( int i = 0; i < 6; i++ ) { // num == 6
		int i = paramMode - 2;
		predAng[i*5 + 3] =init_predAng[i*5+3] + 10*d1;
		//predAng[i*5 + 4] =init_predAng[i*5+3] + 50*d2; 
	} // total of 6*5 = 30 params (5 per leg, 6 legs)
	

}

void BugTracker::AddNextRealView(CvMat *inputRealView, int camIdx)
{
	if( NULL !=realViews[camIdx] ) {
		cvReleaseMat( &(realViews[camIdx]) );
	}

	if( ! GetSystemMetrics( SM_REMOTESESSION )  ) { // don't waste time when running on the node
		cvShowImage( "RenderReadBackReal", inputRealView );
		cvWaitKey(1);
	}

	realViews[camIdx] = inputRealView;
	if( 0 == camIdx ) {
		steps++;
		this->bIsConverged = false; // not converged, we just got new data

		// initialize 'current best' estimates from the prediction at previous time
		this->curAng = init_predAng;
		this->curT	 = init_predT;
		this->curR	 = init_predR;	
		
	}
	attempts = 0;
	// make the dists per-camera!
	distBestUnNormalized[camIdx] = 0.0;
	distBest = 1e7;
	paramMode = 0;

	// now perform a segmentation of the input real view.
	// make array(s) of "yes bug" pixel indices, an array for each cam view
	this->pixelIdxReal[camIdx].clear();
	SegmentGrayThreshold(inputRealView, pixelIdxReal[camIdx],120,true);
	
	if( ! GetSystemMetrics( SM_REMOTESESSION )  ) { // don't waste time when running on the node
		cvShowImage( "SegmentReal", inputRealView );
		cvWaitKey(1);
	}

	srand( time(NULL) );
}

void BugTracker::AddNextGenView(CvMat *inputGenView, int camIdx)
{
	if( NULL != genViews[camIdx] ) {
		cvReleaseMat( &(genViews[camIdx]) );
	}
	
	
	
	vector<int> idxPrv;
	if( attempts > 0 ) {
		idxPrv = pixelIdxGenPrev[camIdx];
		this->pixelIdxGen[camIdx].clear();
		this->pixelIdxGenIn[camIdx].clear();
		this->pixelIdxGenOut[camIdx].clear();
	}

	SegmentGrayThreshold(inputGenView, pixelIdxGen[camIdx],110,true);
	vector<int> idxCur = pixelIdxGen[camIdx];
	
	if( attempts > 0 ) {
		vector<int> both; // indices in both previous and current Gen views
		set_union(idxCur.begin(),idxCur.end(),
			idxPrv.begin(), idxPrv.end(), 
			inserter( both, both.begin() ) );
		
		// in the current but not previous -> index came "in"
		vector<int> out1;
		set_difference(both.begin(),both.end(),
			idxPrv.begin(), idxPrv.end(), 
			inserter( out1, out1.begin() ) );
		pixelIdxGenIn[camIdx] = out1;

		// in the previous but not current -> index went "out"
		vector<int> out2;
		set_difference(both.begin(),both.end(),
			 idxCur.begin(),idxCur.end(),
			inserter( out2, out2.begin() ));
		pixelIdxGenOut[camIdx] = out2;
	}

	if( 0 == attempts ) {
		pixelIdxGenIn[camIdx] = idxCur;
		pixelIdxGenOut[camIdx].clear();
	}

	this->segPixCount[camIdx] = idxCur.size();
	

	genViews[camIdx] = inputGenView;
	
	if( ! GetSystemMetrics( SM_REMOTESESSION )  ) { // don't waste time when running on the node
		cvShowImage( "RenderReadBackSynthetic", inputGenView );
		cvWaitKey(1);
	}




}

void BugTracker::SegmentGrayThreshold(CvMat* img, vector<int>& interiorIdx, int threshold, bool updateimg )
{
	interiorIdx.clear();
	int channels = 3; // should get it from the img actually
	int height = img->height;
	int width = img->width;
	int k = 0;
	for( k = 0; k < height*width; k++ ) {
		int i = k / width;
		int j = k % width;
		CvScalar s = cvGet2D( img, i,j );
		float grayval = (s.val[0] + s.val[1] + s.val[2])/3.0;
		if( grayval > threshold ) {
			interiorIdx.push_back(k);
			//interiorIdx.insert(k);
			if( updateimg ) {
				cvSet2D(img,i,j,cvScalar(255.0,255.0,255.0) );
			}
		}
		else if( updateimg )
			cvSet2D(img,i,j,cvScalar(0.0,0.0,0.0) );
	}

	sort( interiorIdx.begin(), interiorIdx.end() );


}